// "Copyright [2021] <Copyright Shengkai Lin>"

#include <memory>
#include "route.h"

using namespace omnetpp;

namespace ai_sim {

void IBClusterRouteAlgorithm::generateRoutes(taskboard& tmp_taskboard, TopoManager * topos) 
{
    for (auto& routing_stage : tmp_taskboard.stage_info) 
    {
    // Loop in different stages

        for (auto& comm_pair : routing_stage.comm_pairs) 
        {
            topos->getDataplaneTopo()->calculateWBSSPathsTo(comm_pair.second);
            auto tmp_node = comm_pair.first;

            int first_id = topos->getWorkersTopo()->getNodeId(tmp_node);
            int second_id = topos->getWorkersTopo()->getNodeId(comm_pair.second);
            int gate_id_A = -1, gate_id_B = -1, gate_id_C = -1, gate_id_D = -1;

            EV << "--------------" << "\n";
            EV << "Src is " <<
                tmp_node->getModule()->getFullPath() << "\n";
            EV << "Dest is " << comm_pair.second->getModule()->
                getFullPath() << std::endl;
            EV << "Dist is " << tmp_node->getDistanceToTarget() << std::endl;
            if (tmp_node->getNumPaths() == 0) 
            {
                throw omnetpp::cRuntimeError("CAN NOT ALLOC ANY ROUTES, CHECK THE TASK");
                return;
            }

            // First, set the output port of the first gpu.
            for (auto& flow : routing_stage.flows_on_gpu[first_id]) 
            {
                if (flow.destIdx == second_id) 
                {
                    flow.outportId = tmp_node->getoutPathtoDest(0)->getLocalGateId();
                }
            }
            // Second, Start alloc routes for switches

            // Record the specfied links
            routing_stage.used_links.push_back(tmp_node->getoutPathtoDest(0));
            routing_stage.used_links.push_back(tmp_node->getinPathfromDest(0));

            gate_id_A = tmp_node->getoutPathtoDest(0)->getRemoteGateId();
            gate_id_D = tmp_node->getinPathfromDest(0)->getRemoteGateId();

            tmp_node = tmp_node->getoutPathtoDest(0)->getRemoteNode();

            while (tmp_node->getDistanceToTarget() > 0) 
            {
                gate_id_B = tmp_node->getoutPathtoDest(0)->getLocalGateId();
                gate_id_C = tmp_node->getinPathfromDest(0)->getLocalGateId();
                RouteInfo tmp_route_f = 
                {
                    .inPort = gate_id_A,
                    .outPort = gate_id_B,
                    .srcIdx = first_id,
                    .destIdx = second_id
                };
                RouteInfo tmp_route_b = 
                {
                    .inPort = gate_id_C,
                    .outPort = gate_id_D,
                    .srcIdx = second_id,
                    .destIdx = first_id
                };
                routing_stage.routes_on_switch[tmp_node->getModuleId()]
                    .push_back(tmp_route_f);
                // Push back depended routes for flows
                for (auto& flow : routing_stage.flows_on_gpu[first_id]) 
                {
                    if (flow.destIdx == second_id) 
                    {
                        flow.dependRoute[tmp_node->getModuleId()].
                            push_back(tmp_route_f);
                    }
                }
                routing_stage.routes_on_switch[tmp_node->getModuleId()]
                    .push_back(tmp_route_b);
                // Push back depended routes for flows
                for (auto& flow : routing_stage.flows_on_gpu[second_id]) 
                {
                    if (flow.destIdx == first_id) 
                    {
                        flow.dependRoute[tmp_node->getModuleId()].
                            push_back(tmp_route_b);
                    }
                }

                EV << "We are in " <<
                    tmp_node->getModule()->getFullPath() << std::endl;
                EV << "Dest is " << comm_pair.second->getModule()->
                    getFullPath() << std::endl;
                EV << tmp_node->getDistanceToTarget() << " hops to go\n";
                routing_stage.used_links.push_back(tmp_node->getoutPathtoDest(0));
                routing_stage.used_links.push_back(tmp_node->getinPathfromDest(0));

                gate_id_A = tmp_node->getoutPathtoDest(0)->getRemoteGateId();
                gate_id_D = tmp_node->getinPathfromDest(0)->getRemoteGateId();

                tmp_node = tmp_node->getoutPathtoDest(0)->getRemoteNode();
            }

            // Finally, update the output_port for the backword flows
            for (auto& flow : routing_stage.flows_on_gpu[second_id]) 
            {
                if (flow.destIdx == first_id) 
                {
                    flow.outportId = gate_id_D;
                }
            }
        }
    }
}

}  // namespace ai_sim
